# 作者: YYJ
# 2025年06月27日21时12分00秒
# 2486249622@qq.com
import random

RED = 0
BLACK = 1


class RBNode:
    def __init__(self, value):
        self.value = value
        self.left = None
        self.right = None
        self.p = None
        self.color = RED


class RBTree:
    def __init__(self):
        self.root = None

    def left_rotate(self, x: RBNode):
        """
        对红黑树中的节点x执行左旋操作。
        左旋转是一种用于维护红黑树平衡性的转换操作。它涉及重新配置树的结构，同时保持二叉搜索树的性质。
        当节点x右子树较重（即它有一个右孩子）时使用此操作，通过围绕x进行左旋转来重新平衡树。
        - x: RBNode，要执行左旋操作的节点。
        return:
        """
        # 将x的右子节点保存为y
        y = x.right
        # 将x的父节点保存为px
        px = x.p
        # 如果x是树的根，则将y设为新的根
        if px is None:
            self.root = y
        # 如果x是其父节点px的左子节点，则将px的左子节点更新为y
        elif px.left == x:
            px.left = y
        # 如果x是其父节点px的右子节点，则将px的右子节点更新为y
        else:
            px.right = y
        # 更新y的父节点为px
        y.p = px
        # 如果y有左子节点，则将其父节点更新为x
        if y.left:
            y.left.p = x
        # 将y的左子节点移动到x的右子节点位置
        x.right = y.left
        # 将x移动到y的左子节点位置
        y.left = x
        # 更新x的父节点为y
        x.p = y

    def right_rotate(self, x):
        """
        对红黑树中的节点x执行右旋操作。
        右旋转是一种用于维护红黑树平衡性的转换操作。它涉及重新配置树的结构，同时保持二叉搜索树的性质。
        当节点x左子树较重（即它有一个左孩子）时使用此操作，通过围绕x进行右旋转来重新平衡树。
        - x: RBNode，要执行右旋操作的节点。
        return:
        """
        # 将x的左子节点保存为y
        y = x.left
        # 将x的父节点保存为px
        px = x.p
        # 如果x是树的根，则将y设为新的根
        if px is None:
            self.root = y
        # 如果x是其父节点px的右子节点，则将px的右子节点更新为y
        elif px.right == x:
            px.right = y
        # 如果x是其父节点px的左子节点，则将px的左子节点更新为y
        else:
            px.left = y
        # 更新y的父节点为px
        y.p = px
        # 如果y有右子节点，则将其父节点更新为x
        if y.right:
            y.right.p = x
        # 将y的右子节点移动到x的左子节点位置
        x.left = y.right
        # 将x移动到y的右子节点位置
        y.right = x
        # 更新x的父节点为y
        x.p = y

    def insert(self, value):
        """
        插入一个新节点到红黑树中，并保持红黑树的性质。
        此函数首先创建一个新的红黑树节点，然后通过比较节点值找到合适的插入位置。
        插入后，通过调用insert_fixup函数来修复可能违反的红黑树性质。
        : value: 要插入节点的值。
        : return:
        """
        # 创建一个新的红黑树节点
        node = RBNode(value)

        # 如果树为空，将新节点设置为根节点
        if not self.root:
            self.root = node
        else:
            # 从根节点开始，遍历树以找到插入位置
            x = self.root
            while x:
                y = x
                # 根据节点值的大小决定是向左还是向右移动
                if node.value < x.value:
                    x = x.left
                else:
                    x = x.right

            # 确定新节点的插入位置
            if node.value < y.value:
                y.left = node
            else:
                y.right = node

            # 设置新节点的父节点
            node.p = y

        # 插入修复，以保持红黑树的性质
        self.insert_fixup(node)

    def insert_fixup(self, node):
        parent = node.p
        while parent and parent.color == RED:
            g_parent = parent.p
            if g_parent is None:
                break
            if g_parent.left is parent:
                # 父节点是爷节点的左孩子
                uncle = g_parent.right
                if uncle and uncle.color == RED:
                    # 叔节点是红色的
                    g_parent.color = RED
                    uncle.color = BLACK
                    parent.color = BLACK
                    node = g_parent
                    parent = node.p
                else:
                    if parent.right is node:
                        # 节点是父节点的右孩子
                        self.left_rotate(parent)
                    # 节点是父节点的左孩子
                    self.right_rotate(parent)
                    g_parent.color = RED
                    parent.color = BLACK

            else:
                uncle = parent.left
                if uncle and uncle.color == RED:
                    # 叔节点是红色的
                    parent.color = uncle.color = BLACK
                    g_parent.color = RED
                    node = g_parent
                    parent = node.p
                else:
                    if parent.left is node:
                        # 节点是父节点的左孩子
                        self.right_rotate(parent)
                    # 节点是父节点的右孩子
                    self.left_rotate(parent)
                    g_parent.color = RED
                    parent.color = BLACK
        self.root.color = BLACK

    def inorder(self, node: RBNode):
        if node:
            self.inorder(node.left)
            print(node.value, end=' ')
            self.inorder(node.right)

    def rbtree_print(self, node, key, direction):
        if node:
            if direction == 0:
                print('%2d(B) is root' % node.value)
            else:
                print("%2d(%s) is %2d's %6s child" % (
                    node.value, ("B" if node.color == 1 else "R"), key, ("right" if direction ==
                    1 else "left")))

            self.rbtree_print(node.left, node.value, -1)
            self.rbtree_print(node.right, node.value, 1)

if __name__ == '__main__':
    number_list = (7, 4, 1, 8, 5, 2, 9, 6, 3)
    tree = RBTree()
    for i in number_list:
        tree.insert(i)
    # tree.inorder(tree.root)
    tree.rbtree_print(tree.root, 0, 0)